import cv2
import numpy as np
from ultralytics import YOLO

# 加载模型
model = YOLO("yolov8n.pt")



def is_inside_fence(point, fence_polygon):
    """判断点是否在电子围栏内"""
    return cv2.pointPolygonTest(fence_polygon, point, False) >= 0

# 打开视频
# 设置电子围栏区域（多边形）
# 示例：一个四边形区域（你可以自己改坐标）

cap = cv2.VideoCapture("video4.mp4")
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
box_w, box_h = 400, 400  # 你可以自己改大改小
cx, cy = frame_width // 2, frame_height // 2

# 自动构造中心围栏
fence_polygon = np.array([
    [cx - box_w // 2, cy - box_h // 2],
    [cx + box_w // 2, cy - box_h // 2],
    [cx + box_w // 2, cy + box_h // 2],
    [cx - box_w // 2, cy + box_h // 2]
], dtype=np.int32)

fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
out = cv2.VideoWriter("video4_fence_alert.mp4", cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 检测人体
    results = model(frame)

    # 画围栏
    cv2.polylines(frame, [fence_polygon], isClosed=True, color=(255, 255, 0), thickness=2)

    # 检测每个目标
    for r in results:
        for box in r.boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            if cls_id == 0 and conf > 0.5:  # 人体
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                cx, cy = (x1 + x2) // 2, (y1 + y2) // 2

                inside = is_inside_fence((cx, cy), fence_polygon)
                color = (0, 0, 255) if inside else (0, 255, 0)
                label = "ALERT: Fence Breach!" if inside else f"Person ({conf:.2f})"

                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                cv2.circle(frame, (cx, cy), 4, color, -1)
                cv2.putText(frame, label, (x1, y1 - 8),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    out.write(frame)
    cv2.imshow("Fence Detection", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()
